#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist, TransformStamped
from nav_msgs.msg import Odometry
from tf2_ros import TransformBroadcaster
import math
from M80106 import MotorDriver
import time
import asyncio
from concurrent.futures import ThreadPoolExecutor
from functools import partial
from threading import Lock
from tf2_ros.static_transform_broadcaster import StaticTransformBroadcaster

class MotorController(Node):
    def __init__(self):
        super().__init__('motor_controller')
        # 添加互斥锁初始化到最前面，确保即使后续代码失败，锁也已经被创建
        self.motor_lock = Lock()

        # 机器人物理参数
        self.wheel_radius = 0.203 / 2  # 轮子直径203mm转换为半径(m)
        self.wheel_separation = 0.5866  # 轮距586.6mm转换为米
        self.gear_ratio = 6.33  # 电机减速比
        
        # 初始化电机驱动
        try:
            self.motor_driver = MotorDriver()
            # 检查通信线程状态
            if self.motor_driver._running and self.motor_driver._comm_thread.is_alive():
                self.get_logger().info('电机通信线程正常运行')
            else:
                self.get_logger().error('电机通信线程未启动')
                # 尝试重新启动通信
                self.motor_driver.start_comm()
                time.sleep(0.1)  # 等待线程启动
                if self.motor_driver._running and self.motor_driver._comm_thread.is_alive():
                    self.get_logger().info('电机通信线程重启成功')
                else:
                    self.get_logger().error('电机通信线程重启失败')
                    return
        except Exception as e:
            self.get_logger().error(f'电机驱动初始化失败: {str(e)}')
            return

        # 位姿状态
        self.x = 0.0
        self.y = 0.0
        self.theta = 0.0
        
        # 创建发布器和TF广播器
        self.odom_pub = self.create_publisher(Odometry, 'odom', 10)
        self.tf_broadcaster = TransformBroadcaster(self)
        
        # 订阅cmd_vel
        self.subscription = self.create_subscription(
            Twist, 
            '/cmd_vel', 
            self.twist_callback,  # 使用twist_callback作为唯一的回调
            10
        )
        self.timer = self.create_timer(0.2, self.update_odom)
        
        self.last_linear_x = 0.0
        self.last_angular_z = 0.0
        
        # 添加电机位置记录
        self.last_motor_pos = 0.0  # 记录上次电机位置
        self.last_time = self.get_clock().now()  # 记录上次更新时间
        
        # 状态数据
        self.left_status = None
        self.right_status = None
        
        # 当前控制指令
        self.current_cmd = {
            'torque': 0.0,
            'mode': 0x00  # 默认停止模式
        }
        
        # 修改定时器频率为100Hz
        self.create_timer(0.01, self.motor_comm_loop)  # 100Hz的通信循环
        
        # 添加默认状态记录1
        self.last_update_time = self.get_clock().now()
        self.is_moving = False
        self.no_motion_count = 0

        self.get_logger().set_level(rclpy.logging.LoggingSeverity.DEBUG)

        # 修改：使用专门的静态TF广播器
        self.static_broadcaster = StaticTransformBroadcaster(self)
        
        # 立即发布一次静态TF，而不是使用定时器
        self.publish_static_tf()

        # 初始化里程计发布器
        self.odom_pub = self.create_publisher(Odometry, 'odom', 10)
        
        # 初始化变量
        self.x = 0.0
        self.y = 0.0
        self.theta = 0.0
        self.last_time = self.get_clock().now()
        
        # 添加调试定时器���秒打印一次状态
        self.create_timer(1.0, self.debug_timer_callback)

        # 降低更新阈值
        self.SPEED_THRESHOLD = 0.01  # 降低到1cm/s
        self.DISTANCE_THRESHOLD = 0.0005  # 降低到0.5mm

        # 添加当前速度状态记录
        self.current_linear_x = 0.0
        self.current_angular_z = 0.0
        
        # 创建10Hz的定时器用于发布TF和里程计
        self.create_timer(0.1, self.publish_timer_callback)  # 10Hz
        
        # 添加左右轮位置记录
        self.last_left_pos = 0.0
        self.last_right_pos = 0.0
        
        # 降低位移检测阈值
        self.DISTANCE_THRESHOLD = 0.0001  # 0.1mm
        
        # 确保motor_comm_loop的频率足够高
        self.create_timer(0.01, self.motor_comm_loop)  # 100Hz
        
        # 创建更高频率的里程计更新定时器（50Hz）
        self.create_timer(0.02, self.update_odom)

        # 添加PID控制器参数
        self.Kp = 0.8  # 比例系数
        self.Ki = 0.2  # 积分系数
        self.Kd = 0.05  # 微分系数
        
        # 每个电机的PID状态
        self.pid_states = {
            0: {'error_sum': 0.0, 'last_error': 0.0, 'target_speed': 0.0},
            1: {'error_sum': 0.0, 'last_error': 0.0, 'target_speed': 0.0}
        }
        
        # 添加��度控制循环（50Hz）
        self.create_timer(0.02, self.velocity_control_loop)
        
        # 添加积分限幅
        self.max_integral = 2.0

    def motor_comm_loop(self):
        """100Hz的电机通信循环"""
        with self.motor_lock:
            try:
                # 检查通信线程状态
                if not (self.motor_driver._running and self.motor_driver._comm_thread.is_alive()):
                    self.get_logger().error('通信线程已停止，尝试重启')
                    self.motor_driver.start_comm()
                    return
                
                # 获取电机状态
                left_queue = self.motor_driver._motor_status[0]
                right_queue = self.motor_driver._motor_status[1]
                
                if left_queue and right_queue:
                    self.left_status = left_queue[-1]
                    self.right_status = right_queue[-1]
                    # 只在调试模式下打印电机状态
                    if self.get_logger().get_effective_level() <= rclpy.logging.LoggingSeverity.DEBUG:
                        self.get_logger().debug(
                            f'电机编码器 - 左轮: pos={self.left_status["position"]:.4f}, '
                            f'spd={self.left_status["speed"]:.4f} | '
                            f'右轮: pos={self.right_status["position"]:.4f}, '
                            f'spd={self.right_status["speed"]:.4f}'
                        )
                    
            except Exception as e:
                self.get_logger().error(f'电机通信失败: {str(e)}')
    
    def velocity_control_loop(self):
        """闭环速度控制"""
        try:
            with self.motor_lock:
                for motor_id in [0, 1]:
                    # 获取当前速度并添加日志
                    status = self.motor_driver.get_motor_status(motor_id)
                    self.get_logger().debug(f'电机{motor_id}状态: {status}')  # 新增
                    
                    if not status:
                        self.get_logger().warn(f'电机{motor_id}未获取到状态')  # 新增
                        continue
                        
                    current_speed = status['speed']
                    target_speed = self.pid_states[motor_id]['target_speed']
                    
                    # 添加目标速度和当前速度的日志
                    self.get_logger().debug(
                        f'电机{motor_id} - 目标速度:{target_speed:.4f}, '
                        f'当前速度:{current_speed:.4f}'
                    )  # 新增
                    
                    # 计算误差
                    error = target_speed - current_speed
                    
                    # 更新积分项（带限幅）
                    self.pid_states[motor_id]['error_sum'] += error * 0.02  # dt = 0.02s
                    self.pid_states[motor_id]['error_sum'] = max(
                        min(self.pid_states[motor_id]['error_sum'], 
                            self.max_integral), 
                        -self.max_integral
                    )
                    
                    # 计算微分项
                    error_diff = (error - self.pid_states[motor_id]['last_error']) / 0.02
                    self.pid_states[motor_id]['last_error'] = error
                    
                    # PID控制器输出
                    output = (self.Kp * error + 
                            self.Ki * self.pid_states[motor_id]['error_sum'] + 
                            self.Kd * error_diff)
                    
                    # 限制输出范围
                    output = max(min(output, 1.0), -1.0)
                    
                    # 在实际发送命令前添加日志
                    self.get_logger().debug(f'电机{motor_id} - 发送控制输出:{output:.4f}')  # 新增
                    self.motor_driver.set_motor_speed(motor_id, output)
                    
                    # 记录调试信息
                    if abs(error) > 0.1:  # 只在误差较大时记录
                        self.get_logger().debug(
                            f'电机{motor_id} - 目标:{target_speed:.2f} '
                            f'当前:{current_speed:.2f} '
                            f'误差:{error:.2f} '
                            f'输出:{output:.2f}'
                        )
                        
        except Exception as e:
            self.get_logger().error(f'速度控制失败: {str(e)}')
            import traceback  # 新增
            self.get_logger().error(traceback.format_exc())  # 新增

    def twist_callback(self, msg):
        """速度回调函数"""
        with self.motor_lock:
            try:
                # 添加更详细的速度转换日志
                self.get_logger().debug(
                    f'速度指令转换:\n'
                    f'输入 - linear:{msg.linear.x:.4f}, angular:{msg.angular.z:.4f}'
                )  # 新增
                
                # 记录接收到的命令
                self.get_logger().info(f'收到速度命令: linear={msg.linear.x}, angular={msg.angular.z}')
                
                # 更新当前速度状态
                self.current_linear_x = msg.linear.x
                self.current_angular_z = msg.angular.z
                
                # 运动学逆解计算目标速度
                left_wheel_vel = (2.0 * (-msg.linear.x) - msg.angular.z * self.wheel_separation) / (2.0 * self.wheel_radius)
                right_wheel_vel = (2.0 * (-msg.linear.x) + msg.angular.z * self.wheel_separation) / (2.0 * self.wheel_radius)
                
                # 添加轮速计算结果日志
                self.get_logger().debug(
                    f'轮速计算结果:\n'
                    f'左轮目标速度:{left_wheel_vel:.4f} rad/s\n'
                    f'右轮目标速度:{right_wheel_vel:.4f} rad/s'
                )  # 新增
                
                # 更新PID控制器的目标速度
                self.pid_states[0]['target_speed'] = left_wheel_vel
                self.pid_states[1]['target_speed'] = right_wheel_vel
                
                # 如果是零速度命令，清除积分项
                if abs(msg.linear.x) < 0.001 and abs(msg.angular.z) < 0.001:
                    for motor_id in [0, 1]:
                        self.pid_states[motor_id]['error_sum'] = 0.0
                        self.pid_states[motor_id]['target_speed'] = 0.0
                        self.motor_driver.set_motor_speed(motor_id, 0)
                
            except Exception as e:
                self.get_logger().error(f'电机控制失败: {str(e)}')
                import traceback  # 新增
                self.get_logger().error(traceback.format_exc())  # 新增

    def update_odom(self):
        """直接显示电机原始数据"""
        try:
            # 获取电机状态
            left_status = self.motor_driver.get_motor_status(0)
            right_status = self.motor_driver.get_motor_status(1)
            
            if left_status and right_status:
                # 打印完整的原始状态据
                self.get_logger().info(
                    f'\n左轮原始数据:\n'
                    f'  位置(raw): {left_status["position"]}\n'
                    f'  速度(raw): {left_status["speed"]}\n'
                    f'  模式: {left_status["mode"]}\n'
                    f'  力矩: {left_status["torque"]}\n'
                    f'  温度: {left_status["temperature"]}\n'
                    f'\n右轮原始数据:\n'
                    f'  位置(raw): {right_status["position"]}\n'
                    f'  速度(raw): {right_status["speed"]}\n'
                    f'  模式: {right_status["mode"]}\n'
                    f'  力矩: {right_status["torque"]}\n'
                    f'  温度: {right_status["temperature"]}'
                )
                
        except Exception as e:
            self.get_logger().error(f'获取电机数据失败: {str(e)}')

    def publish_tf_and_odom(self, linear_x, angular_z):
        """发布TF和里程计消息"""
        try:
            # 删除或注释掉调试日志
            # self.get_logger().debug('发布TF: odom -> base_link')
            
            t = TransformStamped()
            t.header.stamp = self.get_clock().now().to_msg()
            t.header.frame_id = 'odom'
            t.child_frame_id = 'base_link'
            t.transform.translation.x = float(self.x)
            t.transform.translation.y = float(self.y)
            t.transform.translation.z = 0.0
            
            # 四元数转换
            t.transform.rotation.x = 0.0
            t.transform.rotation.y = 0.0
            t.transform.rotation.z = float(math.sin(self.theta / 2.0))
            t.transform.rotation.w = float(math.cos(self.theta / 2.0))
            
            # 发布TF
            self.tf_broadcaster.sendTransform(t)
            
            # 创建并发布里程计消息
            odom = Odometry()
            odom.header.frame_id = 'odom'
            odom.child_frame_id = 'base_link'
            odom.header.stamp = self.get_clock().now().to_msg()
            
            # 设置置
            odom.pose.pose.position.x = self.x
            odom.pose.pose.position.y = self.y
            odom.pose.pose.position.z = 0.0
            odom.pose.pose.orientation = t.transform.rotation
            
            # 设置协方差（添加这部分以提高稳定性）
            odom.pose.covariance[0] = 0.1  # x位置方差
            odom.pose.covariance[7] = 0.1  # y位置方差
            odom.pose.covariance[35] = 0.1 # 旋转方差
            
            # 设置速度
            odom.twist.twist.linear.x = linear_x
            odom.twist.twist.angular.z = angular_z
            
            # 速度协方差
            odom.twist.covariance[0] = 0.1  # 线速度方差
            odom.twist.covariance[35] = 0.1 # 角速度方差
            
            # 发布里程计消息
            self.odom_pub.publish(odom)
            
        except Exception as e:
            self.get_logger().error(f'发布里程计失败: {str(e)}')

    def publish_static_tf(self):
        """发布静态TF: map -> odom"""
        t = TransformStamped()
        t.header.stamp = self.get_clock().now().to_msg()
        t.header.frame_id = 'map'
        t.child_frame_id = 'odom'
        
        t.transform.translation.x = 0.0
        t.transform.translation.y = 0.0
        t.transform.translation.z = 0.0
        t.transform.rotation.x = 0.0
        t.transform.rotation.y = 0.0
        t.transform.rotation.z = 0.0
        t.transform.rotation.w = 1.0
        
        # 使用静态TF广播器发送变换
        self.static_broadcaster.sendTransform(t)

    def destroy_node(self):
        """清理资源"""
        with self.motor_lock:
            self.motor_driver.stop_motor(0)
            self.motor_driver.stop_motor(1)
            self.motor_driver.close()
        return super().destroy_node()

    def debug_timer_callback(self):
        """调试定时器回调"""
        self.get_logger().info(f'当前位置: x={self.x:.2f}, y={self.y:.2f}, theta={self.theta:.2f}')

    def publish_timer_callback(self):
        """固定频率发布TF和里程计"""
        try:
            current_time = self.get_clock().now()
            dt = (current_time - self.last_time).nanoseconds / 1e9
            
            # 使用当前速度更新位置
            delta_x = self.current_linear_x * math.cos(self.theta) * dt
            delta_y = self.current_linear_x * math.sin(self.theta) * dt
            delta_theta = self.current_angular_z * dt
            
            # 更���位置
            self.x += delta_x
            self.y += delta_y
            self.theta += delta_theta
            
            # 更新间戳
            self.last_time = current_time
            
            # 发布TF和里程计
            self.publish_tf_and_odom(self.current_linear_x, self.current_angular_z)
            
        except Exception as e:
            self.get_logger().error(f'定时发布失败: {str(e)}')

def main():
    rclpy.init()
    motor_controller = MotorController()
    try:
        rclpy.spin(motor_controller)
    except KeyboardInterrupt:
        pass
    finally:
        motor_controller.destroy_node()
        try:
            rclpy.shutdown()
        except Exception as e:
            pass  # 忽略shutdown错误

if __name__ == '__main__':
    main()
